
import rpyc

from rpyc.utils.server import ThreadedServer

from sqlalchemy import create_engine, Binary, Column, Integer, String, Boolean, \
                            MetaData, Table, PickleType, Float, TIMESTAMP, ForeignKey

from sqlalchemy.orm import sessionmaker, query, mapper, join

from sqlalchemy.sql import select, and_, or_, func

from db_tables import Base, Schema, ScriptSettings, PopulationSettings, ColumnSettings

import os, sys, time, traceback, types, cPickle, pickle
from threading import Thread

""" MODIFIES DEFAULT BEHAVIOUR OF SQLALCHEMY """
from sqlalchemy.databases.sqlite import ischema_names, SLFloat
from sqlalchemy.types import type_map
ischema_names['FLOAT'] = SLFloat
type_map[float] = Float

temp_dir = os.path.abspath(__file__)
for i in range(4):
    temp_dir = os.path.dirname(temp_dir)
sys.path.append(temp_dir)

from dexen.libs import constants
from dexen.system.libs import _constants 

""" GLOBALS """
python_to_sql_type = {}
python_to_sql_type[constants.BOOLEAN] = Boolean
python_to_sql_type[constants.STRING] = String
python_to_sql_type[constants.FLOAT] = Float
python_to_sql_type[constants.OBJ] = Binary                
python_to_sql_type[constants.INT] = Integer

PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL
 
class SettingsAPIService(rpyc.Service):
    
    def on_connect(self):
        print "Somebody connected in SettingsAPIService"
        self.settings_session = SettingsSession()
        self.scripts_session = ScriptsSession()
    
    def on_disconnect(self):
        pass
    
    def getPopulationObj(self):
        session = self.settings_session
        schema_name = self.schema_name
        pop_name = self.pop_name
        
        pop_settings = session.query(PopulationSettings).\
                            filter(PopulationSettings.name == pop_name).\
                            join(Schema).\
                            filter(Schema.name == schema_name).first()
        return pop_settings        

    def exposed_setRPYCObj(self, schema_name, pop_name):
        self.schema_name = schema_name
        self.pop_name = pop_name
        self.pop_settings = self.getPopulationObj()
        
    """ Later might convert to join """
    def exposed_setImage(self, image_name):
        session = self.settings_session
        pop_settings = self.getPopulationObj()
        pop_settings.image_name = image_name        
        session.commit()

    def exposed_setModel(self, model_name):
        session = self.settings_session
        pop_settings = self.getPopulationObj()
        pop_settings.model_name = model_name        
        session.commit()
    
    def exposed_setMaxIndividuals(self, max_individuals):
        session = self.settings_session
        pop_settings = self.getPopulationObj()
        pop_settings.max_individuals = max_individuals            
        session.commit()          
        
    def exposed_setObjectives(self, raw_attributes):
        session = self.settings_session
        attributes = cPickle.loads(raw_attributes)
        pop_settings = self.getPopulationObj()
        pop_settings.objectives = attributes                
        session.commit()
    
    def exposed_setScript(self, script_name, file_name, function_name):
        session = self.settings_session
        pop_settings = self.getPopulationObj()
        script_settings = ScriptSettings(script_name, file_name, function_name)
        script_settings.population_settings = pop_settings
        session.add(script_settings)            
        session.commit()
       
    def exposed_setDataStructure(self, raw_attributes):
        try:
            session = self.settings_session
            attributes = cPickle.loads(raw_attributes)
            pop_settings = self.getPopulationObj()
            for pair in attributes.items():
                col_settings = ColumnSettings(pair[0], pair[1])
                col_settings.population_settings = pop_settings
                session.add(col_settings)
            session.commit()
        except:
            tb = sys.exc_traceback
            print "Exception in setDataStrucutre at line ", traceback.tb_lineno(tb), "\n"
 

class ClientAPIService(rpyc.Service):
    def on_connect(self):
        print "Somebody connected in ClientAPIService"
        self.settings_session = SettingsSession()
        self.scripts_session = ScriptsSession()
    
    def on_disconnect(self):
        pass
    
    def exposed_setRPYCObj(self):
        pass    
    
    def getPopulationObj(self, schema_name, pop_name):
        session = self.settings_session
        pop_settings = session.query(PopulationSettings).\
                            filter(PopulationSettings.name == pop_name).\
                            join(Schema).\
                            filter(Schema.name == schema_name).first()
        return pop_settings       
    
    def exposed_setIndividualTable(self, schema_name, pop_name):
        session = self.settings_session 
        metadata = MetaData()
        metadata.bind = scripts_engine
        
        ind_table_name = schema_name + "_" + pop_name
        
        ind_table = Table(ind_table_name, metadata,
                          Column('id', Integer, primary_key=True)
                    )
        
        col_pairs = session.query(ColumnSettings.column_name, ColumnSettings.column_type).\
                        join(PopulationSettings).\
                        filter(PopulationSettings.name == pop_name).\
                        join(Schema).\
                        filter(Schema.name == schema_name).all()
                                                      
        """ ADDING DYNAMIC DATASTRUCTURE COLUMNS HERE """
        for pair in col_pairs:
            ind_table.append_column(Column(pair[0], python_to_sql_type[pair[1]]))
        
        scripts = session.query(ScriptSettings).join(PopulationSettings).\
                        filter(PopulationSettings.name == pop_name).\
                        join(Schema).\
                        filter(Schema.name == schema_name).all()
        
        """ TIMESTAMPS FOR THE SCRIPTS """
        for script in scripts:
            ind_table.append_column(Column(script.name, Float))
            
        ind_table.create(bind=scripts_engine)
        session.commit()  
            
    def exposed_setSchema(self, schema_name, zip_data, last_modified_time):
        session = self.settings_session
        schema = Schema(schema_name, zip_data, last_modified_time)
        session.add(schema)
        session.commit()
        
    def exposed_getSchemas(self):
        session = self.settings_session
        schemas = []
        for schema in session.query(Schema.name, Schema.file_name, Schema.zip_data, Schema.last_modified_time):
            schemas.append(schema)
        session.commit()
        return schemas 
    
    def exposed_getSchema(self, schema_name):
        session = self.settings_session
        schema = session.query(Schema.zip_data, Schema.last_modified_time).\
                        filter(Schema.name == schema_name).first()
        session.commit()
        return schema
    
    def exposed_getSchemaTimestamp(self, schema_name):
        session = self.settings_session
        schema_timestamp = session.query(Schema.last_modified_time).\
                        filter(Schema.name == schema_name).scalar()
        session.commit()
        return schema_timestamp
    
    def exposed_getSchemaNames(self):
        session = self.settings_session
        schema_names = []
        for schema_name in session.query(Schema.name).order_by(Schema.name):
            schema_names.append(schema_name[0])
        session.commit()
        return schema_names
    
    def exposed_stopPopulation(self, schema_name, pop_name):
        session = self.settings_session
        pop_settings = self.getPopulationObj(schema_name, pop_name)
        pop_settings.is_stopped = True
        session.commit()    
    
    def exposed_setPopulation(self, schema_name, pop_name):
        session = self.settings_session
        pop_settings = PopulationSettings(pop_name)
        schema = session.query(Schema).filter(Schema.name == schema_name).first()
        pop_settings.schema = schema
        session.add(pop_settings)
        session.commit()

    
    def exposed_getPopulationNames(self, schema_name):
        session = self.settings_session
        schema = session.query(Schema).filter(Schema.name == schema_name).first()
        pop_names = []
        for pop_settings in schema.populations:
            pop_names.append(pop_settings.name)
        session.commit()
        return pop_names   

    """ Is pop name unique globally ?"""
    def exposed_getScriptFileNames(self, schema_name, pop_name):
        session = self.settings_session
        script_settings = session.query(ScriptSettings).join(PopulationSettings).\
                                join(Schema).\
                                filter(Schema.name == schema_name).\
                                filter(PopulationSettings.name == pop_name).all()
        script_file_names = []
        #print script_settings        
        for script in script_settings:
            script_file_names.append(script.file_name)
        session.commit()
        return script_file_names
    
    def exposed_getScriptFunctionNames(self, schema_name, pop_name, file_name):
        session = self.settings_session
        script_settings = session.query(ScriptSettings).\
                                join(PopulationSettings).\
                                join(Schema).\
                                filter(Schema.name == schema_name).\
                                filter(PopulationSettings.name == pop_name).\
                                filter(ScriptSettings.file_name == file_name).all()
        
        script_function_names = []
        for script in script_settings:
            script_function_names.append(script.function_name)
        session.commit()
        return script_function_names
    
    def exposed_getScriptFileFunctionNames(self, schema_name, pop_name, script_name):
        session = self.settings_session        
        file_name, function_name = session.query(ScriptSettings.file_name, ScriptSettings.function_name).\
                                    join(PopulationSettings).\
                                    join(Schema).\
                                    filter(Schema.name == schema_name).\
                                    filter(PopulationSettings.name == pop_name).\
                                    filter(ScriptSettings.name == script_name).first()
        session.commit()                                              
        return file_name, function_name                             
    
    def exposed_getScriptNames(self, schema_name, pop_name):
        session = self.settings_session
        script_settings = session.query(ScriptSettings).\
                                    join(PopulationSettings).\
                                    join(Schema).\
                                    filter(Schema.name == schema_name).\
                                    filter(PopulationSettings.name == pop_name).all()
        script_names = []
        
        for script in script_settings:
            script_names.append(script.name)
        session.commit()
        return script_names
    

class ManagerAPIService(rpyc.Service):
    
    def on_connect(self):
        print "Somebody connected in ManagerAPIService"
        self.settings_session = SettingsSession()
        self.scripts_session = ScriptsSession()
    
    def on_disconnect(self):
        pass

    def exposed_setRPYCObj(self, schema_name, pop_name):
        self.schema_name = schema_name
        self.pop_name = pop_name
    
    def getPopulationObj(self):
        session = self.settings_session
        schema_name = self.schema_name
        pop_name = self.pop_name
        
        pop_settings = session.query(PopulationSettings).\
                            filter(PopulationSettings.name == pop_name).\
                            join(Schema).\
                            filter(Schema.name == schema_name).first()
        return pop_settings            
    
    def exposed_getFitnessNames(self):
        session = self.settings_session
        pop_settings = self.getPopulationObj()
        return [fitness[0] for fitness in pop_settings.fitness_info]      
    
    def exposed_hasScriptInput(self):
        session = self.settings_session
        schema_name = self.schema_name
        pop_name = self.pop_name
        script_select = session.query(ScriptSettings.cond_select).\
                            join(PopulationSettings).\
                            join(Schema).\
                            filter(PopulationSettings.name == pop_name).\
                            filter(Schema.name == schema_name).\
                            filter(ScriptSettings.name == script_name).scalar()
        
        
        if script_select == None:
            return False
        return True

    def exposed_hasPopulationStopped(self):
        pop_settings = self.getPopulationObj()        
        return pop_settings.is_stopped


class ScriptAPIService(rpyc.Service):
    TIME_OUT = 600 #IN TERMS OF SECONDS 
    def on_connect(self):
        print "Somebody connected in ScriptAPIService"
        self.settings_session = SettingsSession()
        self.scripts_session = ScriptsSession()
           
    def on_disconnect(self):
        pass  
    
    def exposed_setRPYCObj(self, schema_name, pop_name, script_name):
        self.schema_name = schema_name
        self.pop_name = pop_name
        self.script_name = script_name
        self.pop_settings = self.getPopulationObj()
        self.max_individuals = self.pop_settings.max_individuals
        self.script_settings = self.getScriptObj()
        self.data_structure = self.getDataStructure()
        self.ind_table = self.loadIndividualTable()    
        
    def getPopulationObj(self):
        session = self.settings_session
        schema_name = self.schema_name
        pop_name = self.pop_name
        
        pop_settings = session.query(PopulationSettings).\
                            filter(PopulationSettings.name == pop_name).\
                            join(Schema).\
                            filter(Schema.name == schema_name).first()
        return pop_settings                

    def getScriptObj(self):
        session = self.settings_session
        schema_name = self.schema_name
        pop_name = self.pop_name
        script_name = self.script_name
        
        script_settings = session.query(ScriptSettings).\
                            join(PopulationSettings).\
                            join(Schema).\
                            filter(ScriptSettings.name == script_name).\
                            filter(PopulationSettings.name == pop_name).\
                            filter(Schema.name == schema_name).first()
                                    
        return script_settings

    def getDataStructure(self):
        session = self.settings_session
        schema_name = self.schema_name
        pop_name = self.pop_name
        data_structure = {}
        
        col_pairs = session.query(ColumnSettings.column_name, ColumnSettings.column_type).\
                            join(PopulationSettings).\
                            filter(PopulationSettings.name == pop_name).\
                            join(Schema).\
                            filter(Schema.name == schema_name).all()      
           
        for pair in col_pairs:
            data_structure[pair[0]] = pair[1]
        session.commit()
        return data_structure
                
    def getScriptId(self, script_name, pop_name, schema_name):
        session = self.settings_session
        script_id = session.query(ScriptSettings.id).\
                            join(PopulationSettings).\
                            join(Schema).\
                            filter(ScriptSettings.name == script_name).\
                            filter(PopulationSettings.name == pop_name).\
                            filter(SchemaSettings.name == schema_name).scalar()
        return script_id

    def loadIndividualTable(self):
        schema_name = self.schema_name
        pop_name = self.pop_name
        metadata = MetaData()
        metadata.bind = scripts_engine
        ind_table = Table(schema_name + "_" + pop_name, metadata, autoload=True)
        return ind_table

    def countIndividuals(self):
        session = self.scripts_session
        ind_table = self.ind_table
        no_individuals = session.execute(select([func.count(ind_table.c.id)])).scalar()
        #session.commit()
        return no_individuals
    
    def updateIndividual(self, individual):
        session = self.scripts_session
        script_name = self.script_name
        ind_table = self.ind_table
        
        while True:
            try:
                """ This is to reset the EXCLUSIVE_USE """
                individual[script_name] = None
                session.execute(ind_table.update().where(ind_table.c.id == individual['id']), individual)
                break
            except:
                print script_name, ": exception in update individual"
                print sys.exc_info()
                session.rollback()
        
    def insertIndividual(self, individual):
        session = self.scripts_session
        script_name = self.script_name
        ind_table = self.ind_table
        
        #print ind_table.c, "***************"
        #time.sleep(20)
        while True:
            try:    
                session.execute(ind_table.insert(), individual)
                """ I don't have the id of the individual to be inserted that is why I use a separate function """
                break
            except:
                print script_name, ": exception in insert individual at line", sys.exc_traceback.tb_lineno
                print sys.exc_info()[:-1]
                session.rollback()

    def pickleIndividual(self, individual):
        data_structure = self.data_structure
        for key in individual.keys():
            if key in data_structure.keys():
                if data_structure[key] == constants.OBJ:
                    if individual[key] != None:
                        individual[key] = cPickle.dumps(individual[key], PICKLE_PROTOCOL)
                            
    def exposed_setIndividuals(self, raw_individuals):
        print self.script_name, "in set individuals"
        session = self.scripts_session
        ind_table = self.ind_table    
        pop_settings = self.pop_settings
        script_name = self.script_name
        
        cur_time = time.time()
        individuals = cPickle.loads(raw_individuals)
        
        while True:
            try:
                for individual in individuals:
                    self.pickleIndividual(individual) # pickles only the object attributes of an individual            
                    if 'id' in individual:
                        self.updateIndividual(individual)
                    else:
                        self.insertIndividual(individual)                
                no_individuals = self.countIndividuals()
                #print no_individuals, self.max_individuals
                if no_individuals > self.max_individuals:               
                    pop_settings.is_stopped = True
                    self.settings_session.commit()   
                print "it took", time.time() - cur_time, "seconds..."    
                session.commit()    
                break
            except:
                print script_name, ": exception in set individuals", "at line", sys.exc_traceback.tb_lineno
                #time.sleep(5) 
                print sys.exc_info()
                session.rollback()
        print script_name, "in set individuals took ", time.time() - cur_time, "seconds..." 

    def exposed_getIndividuals(self, select, where, count, order_by, exclusive_use):
        print self.script_name, "in get individuals"              
        session = self.scripts_session
        select_tuple = select
        where_tuple = where
        ind_table = self.ind_table
        script_settings = self.script_settings    
        script_name = self.script_name
        exp_backoff = 0.3
        
        while True:
            try:
                query = ind_table.select()

                """ PROCESSING WHERE CLAUSE """
                for item in where_tuple:
                    if type(item[2]) is types.StringType:
                        query = query.where(eval("ind_table.c."+item[0]+item[1]+"\'"+item[2]+"\'"))
                    else:
                        query = query.where(eval("ind_table.c."+item[0]+item[1]+str(item[2])))
                
                cur_time = time.time()
                
                if exclusive_use == True:
                    """ PROCESSING TIMEOUT AND EXCLUSIVE_USE """
                    query = query.where(or_(eval("ind_table.c." + script_name + "==" + str(None)),
                                            eval("ind_table.c." + script_name + "<=" + str(cur_time - self.TIME_OUT))
                                        )
                            )
                            
                """ PROCESSING ORDER BY CLAUSE """       
                if order_by != None:
                    query = query.order_by(eval("ind_table.c." + order_by[0] + "." + order_by[1]))
                else:
                    """ Check Database compatibility """
                    query = query.order_by(func.random())
                
                """ PROCESSING COUNT """
                if count != None:
                    query = query.limit(count)
                
                results = session.execute(query).fetchall()
                
                """                            
                for ind in results:
                    print ind.id
                time.sleep(20)
                """
                
                if count == None or len(results) == count:
                    ind_list = []
                    for row_obj in results[:]:
                        temp = {}
                        temp['id'] = row_obj['id']
                        for select_item in select_tuple:
                            value = row_obj[select_item]
                            if (select_item in self.data_structure.keys()) and (self.data_structure[select_item] == constants.OBJ) and (row_obj[select_item] != None): 
                                temp[select_item] = cPickle.loads(str(row_obj[select_item]))
                            else:
                                temp[select_item] = row_obj[select_item]
                                
                        ind_list.append(temp)
                        if exclusive_use == True:
                            session.execute(ind_table.update().where(eval("ind_table.c." + script_name + " == " + str(row_obj['id']))), 
                                                                                            {script_name : time.time()})                    
                    print "Get Individuals took", time.time() - cur_time, "seconds..."
                    session.commit()                                        
                    return cPickle.dumps(ind_list, PICKLE_PROTOCOL)
                
                pop_settings = self.getPopulationObj()
                """
                    POTENTIAL BOTTLENECK
                """
                if pop_settings.is_stopped == True:
                    return None
                time.sleep(exp_backoff)
                #exp_backoff *= 2
            except:
                print script_name, ": exception in get individuals at line:", sys.exc_traceback.tb_lineno
                print sys.exc_info()[1]
                time.sleep(exp_backoff)
                session.rollback()     
        return None

  
services = [{"service_name": "SettingsAPIService",
             "service_port": _constants._SETTINGS_API_PORT
            },
            {"service_name": "ClientAPIService",
             "service_port": _constants._CLIENT_API_PORT
            },
            {"service_name": "ManagerAPIService", 
             "service_port": _constants._MANAGER_API_PORT
            },
            {"service_name": "ScriptAPIService",
             "service_port": _constants._SCRIPT_API_PORT
            }
]
    
def db_init():
    metadata = MetaData()
    metadata.reflect(scripts_engine)
    metadata.drop_all(scripts_engine)    
    Base.metadata.drop_all(settings_engine)
    Base.metadata.create_all(settings_engine)

def server_starter(service, _port):
    my_server = ThreadedServer(eval(service), port = _port, auto_register=False)
    my_server.start()
    
def start_services():
    for service in services:
        t = Thread(target = server_starter, args = (service["service_name"], service["service_port"]))
        t.start()

if __name__ == "__main__":
    settings_engine = create_engine("sqlite:///dexen_settings.db", echo=False, pool_size=50)
    scripts_engine = create_engine("sqlite:///dexen_scripts.db", echo=False, pool_size=50)
    SettingsSession = sessionmaker(settings_engine)
    ScriptsSession = sessionmaker(scripts_engine) 
    db_init()
    start_services()
    try:
        raw_input()
    except:
        print ""
    finally:
        sys.exit()

